import json

import model

input_text = "查询一下任务编号：TASK212400000503的排产任务表"
cats = "任务编号"
# 读取span标签配置
with open("span/label_config.json", "r") as fileSpan:
    data = json.load(fileSpan)
    match_names = [item['text'] for item in data]
test_texts = {
    "instruction": "你是一个文本实体识别领域的专家，你需要从给定的句子中提取 " + "; ".join(
        match_names) + """ 实体. 以 json 格式输出, 如 [{"entity_text": "TASK202400000501", "entity_label": "任务编号"}] 注意: 1. 输出的每一行都必须是正确的 json 字符串. 2. 找不到任何实体时, 输出"没有找到任何实体".""",
    "input": f"文本类型:{cats};文本:{input_text}",
}

instruction = test_texts['instruction']
input_value = test_texts['input']

messages = [
    {"role": "system", "content": f"{instruction}"},
    {"role": "user", "content": f"{input_value}"}
]

response = model.predict(messages, model.load_tokenizer_model(), model.load_model_ner())
print(response)
